import json
from websocket import create_connection, WebSocketTimeoutException


class WsClient:
    def __init__(self, url):
        self.args = {}
        self.connected = create_connection(url)

    def post(self, msg):
        msg_json = json.dumps(msg)
        self.connected.send(msg_json)

    def param(self, key, val):
        self.args[key] = val

    def receive(self, handler):
        try:
            while True:
                try:
                    result = self.connected.recv()
                    message = json.loads(result)
                    handler(message, self)
                except WebSocketTimeoutException:
                    print("WebSocket连接超时")
        except Exception as e:
            print("WebSocket连接错误: ", e)
        finally:
            self.close()

    def close(self):
        self.connected.close()
        self.connected = None
